package com.xuyanze.es;


import com.alibaba.fastjson.JSON;
import com.xuyanze.es.face.IElasticsearchId;
import com.xuyanze.es.utils.EsWrapperProcessor;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;
import org.springframework.data.elasticsearch.core.ScrolledPage;
import org.springframework.data.elasticsearch.core.aggregation.AggregatedPage;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.util.CollectionUtils;

import javax.annotation.Resource;
import java.lang.reflect.ParameterizedType;
import java.util.ArrayList;
import java.util.List;

/**
 * @author ：lzy
 * @date ：2022-10-08
 * @description：es
 */
public class BaseElasticRepository<T> {

    @Resource
    private ElasticsearchRestTemplate elasticsearchRestTemplate;

    private Class clazz;

    public static final int SCORE_TIME_OUT= 30 * 1000;

    public ElasticsearchRestTemplate getElasticsearchRestTemplate() {
        return elasticsearchRestTemplate;
    }

    public void createIndex(){
        elasticsearchRestTemplate.createIndex(getClazz());
        elasticsearchRestTemplate.putMapping(getClazz());
    }

    public <T> T findById(String id) {
        GetQuery query = GetQuery.getById(id);
        T data = elasticsearchRestTemplate.queryForObject(query, getClazz());
        return data;
    }

    public void deleteById(String id){
        Class<Object> clazz = getClazz();
        elasticsearchRestTemplate.delete(clazz,id);
        elasticsearchRestTemplate.refresh(clazz);
    }

    public <T> T save(T setting) {
        IndexQuery indexQuery = new IndexQueryBuilder()
                .withObject(setting)
                .build();
        elasticsearchRestTemplate.index(indexQuery);
        elasticsearchRestTemplate.refresh(getClazz());
        return setting;
    }

    public <T extends IElasticsearchId> void update(T setting) {
        update(setting,0);
    }

    public <T extends IElasticsearchId> void update(T setting,int retryOnConflict) {
        UpdateRequest updateRequest = new UpdateRequest();
        updateRequest.doc(JSON.toJSONString(setting), XContentType.JSON);
        updateRequest.retryOnConflict(retryOnConflict);

        Class<T> clazz = getClazz();
        UpdateQuery updateQuery = new UpdateQueryBuilder()
                .withId(setting.getId())
                .withUpdateRequest(updateRequest)
                .withClass(clazz)
                .build();
        elasticsearchRestTemplate.update(updateQuery);
        elasticsearchRestTemplate.refresh(clazz);
    }

    public <T extends IElasticsearchId> void saveOrUpdate(T setting) {
        saveOrUpdate(setting,0);
    }

    public <T extends IElasticsearchId> void saveOrUpdate(T setting,int retryOnConflict) {
        UpdateRequest updateRequest = new UpdateRequest();
        updateRequest.doc(JSON.toJSONString(setting), XContentType.JSON);
        updateRequest.retryOnConflict(retryOnConflict);

        Class<T> clazz = getClazz();
        UpdateQuery updateQuery = new UpdateQueryBuilder()
                .withId(setting.getId())
                .withUpdateRequest(updateRequest)
                .withClass(clazz)
                .withDoUpsert(true)
                .build();
        elasticsearchRestTemplate.update(updateQuery);
        elasticsearchRestTemplate.refresh(clazz);
    }

    public <T extends IElasticsearchId> void saveBatch(List<T> datas) {
        if (CollectionUtils.isEmpty(datas)) {
            return;
        }
        List<IndexQuery> queries = new ArrayList<>(datas.size());
        Class<T> clazz = getClazz();
        for (T data : datas) {
            IndexQuery indexQuery = new IndexQueryBuilder()
                    .withObject(data)
                    .build();
            queries.add(indexQuery);
        }
        elasticsearchRestTemplate.bulkIndex(queries);
        elasticsearchRestTemplate.refresh(clazz);
    }

    public <T extends IElasticsearchId> void updateBatch(List<T> datas) {
        updateBatch(datas,0);
    }

    public <T extends IElasticsearchId> void updateBatch(List<T> datas, int retryOnConflict) {
        if (CollectionUtils.isEmpty(datas)) {
            return;
        }
        List<UpdateQuery> queries = new ArrayList<>(datas.size());
        Class<T> clazz = getClazz();
        for (T data : datas) {
            UpdateRequest updateRequest = new UpdateRequest();
            updateRequest.doc(JSON.toJSONString(data), XContentType.JSON);
            updateRequest.retryOnConflict(retryOnConflict);
            UpdateQuery updateQuery = new UpdateQueryBuilder()
                    .withId(data.getId())
                    .withUpdateRequest(updateRequest)
                    .withClass(clazz)
                    .build();
            queries.add(updateQuery);
        }
        elasticsearchRestTemplate.bulkUpdate(queries);
        elasticsearchRestTemplate.refresh(clazz);
    }

    public <T extends IElasticsearchId> void saveOrUpdateBatch(List<T> datas) {
        saveOrUpdateBatch(datas,0);
    }
    public <T extends IElasticsearchId> void saveOrUpdateBatch(List<T> datas, int retryOnConflict) {
        if (CollectionUtils.isEmpty(datas)) {
            return;
        }
        List<UpdateQuery> queries = new ArrayList<>(datas.size());
        Class<T> clazz = getClazz();
        for (T data : datas) {
            UpdateRequest updateRequest = new UpdateRequest();
            updateRequest.doc(JSON.toJSONString(data), XContentType.JSON);
            updateRequest.retryOnConflict(retryOnConflict);
            UpdateQuery updateQuery = new UpdateQueryBuilder()
                    .withId(data.getId())
                    .withUpdateRequest(updateRequest)
                    .withClass(clazz)
                    .withDoUpsert(true)
                    .build();
            queries.add(updateQuery);
        }
        elasticsearchRestTemplate.bulkUpdate(queries);
        elasticsearchRestTemplate.refresh(clazz);
    }

    private <T> Class<T> getClazz(){
        if(clazz == null){
            clazz = (Class<T>) ((ParameterizedType) getClass().getGenericSuperclass()).getActualTypeArguments()[0];
        }
        return clazz;
    }

    private ScrolledPage<T> startScroll(SearchQuery searchQuery){
        return startScroll(SCORE_TIME_OUT, searchQuery);
    }

    private ScrolledPage<T> startScroll(long timeout,SearchQuery searchQuery){
        return elasticsearchRestTemplate.startScroll(timeout, searchQuery, getClazz());
    }

    private ScrolledPage<T> continueScroll(String scrollId){
        return continueScroll(SCORE_TIME_OUT,scrollId);
    }

    private ScrolledPage<T> continueScroll(long timeout,String scrollId){
        return elasticsearchRestTemplate.continueScroll(scrollId,timeout, getClazz());
    }

    public List<T> scrollPage(SearchQuery searchQuery){
        List<T> data = new ArrayList<>();
        ScrolledPage<T> page = startScroll(searchQuery);
        String scrollId = page.getScrollId();
        int totalPages = page.getTotalPages();

        data.addAll(page.getContent());
        for (int i = 1; i < totalPages; i++) {
            ScrolledPage<T> nextPage = continueScroll(scrollId);
            data.addAll(nextPage.getContent());
        }
        elasticsearchRestTemplate.clearScroll(scrollId);
        return data;
    }

    private ScrolledPage<T> getScrollPage(SearchQuery searchQuery){
        ScrolledPage<T> page = startScroll(searchQuery);
        return page;
    }

    public long count(EsLambdaQueryWrapper lambdaQueryWrapper){
        NativeSearchQuery searchQuery = new NativeSearchQueryBuilder()
                .withQuery(EsWrapperProcessor.initQueryBuilder(lambdaQueryWrapper))
                .build();
        return elasticsearchRestTemplate.count(searchQuery,getClazz());
    }

    public void delete(EsLambdaQueryWrapper lambdaQueryWrapper) {
        delete(EsWrapperProcessor.initQueryBuilder(lambdaQueryWrapper),
                2000,
                100L);
    }

    public void delete(BoolQueryBuilder searchQuery, int pageSize, long scrollTimeInMillis){
        DeleteQuery deleteQuery = new DeleteQuery();
        deleteQuery.setQuery(searchQuery);
        deleteQuery.setPageSize(pageSize);
        deleteQuery.setScrollTimeInMillis(scrollTimeInMillis);
        elasticsearchRestTemplate.delete(deleteQuery, getClazz());
    }

    public List<T> queryForList(EsLambdaQueryWrapper lambdaQueryWrapper){
        return elasticsearchRestTemplate.queryForList(
                EsWrapperProcessor.initNativeSearchQuery(lambdaQueryWrapper),
                getClazz()
        );
    }

    public List<T> queryAll(){
        return queryAll(EsLambdaQueryWrapper.wrapper());
    }

    public List<T> queryAll(EsLambdaQueryWrapper lambdaQueryWrapper){
        return queryAll(lambdaQueryWrapper,PageRequest.of(0, 1000));
    }

    public List<T> queryAll(Pageable pageable){
        return queryAll(EsLambdaQueryWrapper.wrapper(),pageable);
    }

    public List<T> queryAll(EsLambdaQueryWrapper lambdaQueryWrapper, Pageable pageable){
        NativeSearchQuery searchQuery = new NativeSearchQueryBuilder()
                .withQuery(EsWrapperProcessor.initQueryBuilder(lambdaQueryWrapper))
                .withPageable(pageable)
                .build();
        return scrollPage(searchQuery);
    }

    public T queryForOnce(EsLambdaQueryWrapper lambdaQueryWrapper){
        List<T> ts = elasticsearchRestTemplate.queryForList(
                EsWrapperProcessor.initNativeSearchQuery(lambdaQueryWrapper),
                getClazz()
        );
        return CollectionUtils.isEmpty(ts) ? null : ts.get(0);
    }


    public Page queryForPage(EsLambdaQueryWrapper lambdaQueryWrapper){
        return elasticsearchRestTemplate.queryForPage(
                EsWrapperProcessor.initNativeSearchQuery(lambdaQueryWrapper),
                getClazz()
        );
    }

    public AggregatedPage<T> queryForAggregatedPage(EsLambdaQueryWrapper lambdaQueryWrapper){
        return elasticsearchRestTemplate.queryForPage(
                EsWrapperProcessor.initNativeSearchQuery(lambdaQueryWrapper),
                getClazz()
        );
    }


}
